from __future__ import division
import tkinter as tk
from tkinter import filedialog
from PIL import Image
import numpy as np



def getDark(input_img, filter, frame):
    """get dark image from the input image"""
    size = input_img.size
    output = []

    for x in range(size[1]):
        temp = []
        for y in range(size[0]):
            temp.append(min(input_img.getpixel((y, x))))

        output.append(temp)

    output = filter2d(output, filter, frame)

    output_img = Image.new('L', size)

    for x in range(size[1]):
        for y in range(size[0]):
            output_img.putpixel((y, x), output[x][y])

    return output_img


def getLight(srcImage, darkImage, cut):
    """get atmospheric light from the picture"""
    size = darkImage.size
    light = []

    for x in range(size[0]):
        for y in range(size[1]):
            light.append(darkImage.getpixel((x, y)))

    light.sort()
    light.reverse()

    threshold = light[int(cut * len(light))]

    atmosphere = {}

    for x in range(size[0]):
        for y in range(size[1]):
            if darkImage.getpixel((x, y)) >= threshold:
                atmosphere.update({(x, y): sum(srcImage.getpixel((x, y))) / 3.0})

    pos = sorted(atmosphere.items(), key=lambda item: item[1], reverse=True)[0][0]

    return srcImage.getpixel(pos)


def getTransmission(input_img, light, omiga):
    """get transmission from the picture"""
    size = input_img.size
    output = []

    for x in range(size[1]):
        temp = []
        for y in range(size[0]):
            temp.append(min(input_img.getpixel((y, x))) / float(min(light)))

        output.append(temp)

    transmission = []

    for x in range(size[1]):
        temp = []
        for y in range(size[0]):
            temp.append(1 - omiga * minimizeFilter(output, (x, y), (10, 10)))

        transmission.append(temp)

    return transmission


def getRadiance(input_img, transmission, light, t0):
    """get radiance from the picture"""
    size = input_img.size
    output = Image.new('RGB', size)

    for x in range(size[1]):
        for y in range(size[0]):
            r, g, b = input_img.getpixel((y, x))

            r = int((r - light[0]) / float(max(t0, transmission[x][y])) + light[0])
            g = int((g - light[1]) / float(max(t0, transmission[x][y])) + light[1])
            b = int((b - light[2]) / float(max(t0, transmission[x][y])) + light[2])

            output.putpixel((y, x), (r, g, b))

    return output


def region_growing(image, seed_point, threshold=30):
    """
    区域生长算法，用于分割天空区域。
    :param image: 输入的灰度图像 (PIL Image 或 numpy array)
    :param seed_point: 种子点 (x, y)
    :param threshold: 相似性阈值，用于决定是否扩展区域
    :return: 分割的二值图像，天空区域为 1，其他区域为 0
    """
    if isinstance(image, Image.Image):
        image = np.array(image)

    height, width = image.shape
    segmented = np.zeros_like(image, dtype=np.uint8)  # 初始化分割结果
    visited = np.zeros_like(image, dtype=np.uint8)  # 记录访问过的像素
    seed_x, seed_y = seed_point
    seed_value = image[seed_y, seed_x]  # 种子点的亮度值

    # 用于区域生长的队列初始化
    queue = [(seed_x, seed_y)]
    segmented[seed_y, seed_x] = 1
    visited[seed_y, seed_x] = 1

    while queue:
        x, y = queue.pop(0)

        # 遍历当前像素的 8 邻域
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1),
                       (-1, -1), (-1, 1), (1, -1), (1, 1)]:
            nx, ny = x + dx, y + dy

            if 0 <= nx < width and 0 <= ny < height and not visited[ny, nx]:
                visited[ny, nx] = 1  # 标记为已访问
                # 判断当前像素是否满足相似性条件
                if abs(int(image[ny, nx]) - int(seed_value)) < threshold:
                    segmented[ny, nx] = 1
                    queue.append((nx, ny))  # 将符合条件的像素加入队列

    return segmented


def process_sky_region(image, sky_mask):
    """
    :param image: 原始图像 (PIL Image)
    :param sky_mask: 天空区域的二值掩码 (numpy array)
    :return: 处理后的透射率图
    """
    # 将天空区域的透射率设置为 1.0
    height, width = sky_mask.shape
    transmission = np.ones((height, width))  # 初始化透射率为 1.0
    for y in range(height):
        for x in range(width):
            if sky_mask[y, x] == 0:  # 非天空区域
                transmission[y, x] = 0.7  # 可根据需要调整非天空区域透射率
            else:  # 天空区域
                transmission[y, x] = 0.5
    return transmission


def ensure(n):
    if n < 0:
        n = 0

    if n > 255:
        n = 255

    return int(n)


def filter2d(data, filter_func, frame):
    """
    简单的二维滤波包装函数，这里只是占位，实际可按需实现更复杂的滤波逻辑
    """
    rows = len(data)
    cols = len(data[0])
    result = [[0] * cols for _ in range(rows)]
    for i in range(rows):
        for j in range(cols):
            result[i][j] = filter_func(data, (j, i), frame)
    return result


def minimizeFilter(data, pos, size):
    """
    最小滤波函数
    :param data: 输入数据
    :param pos: 当前位置 (x, y)
    :param size: 滤波窗口大小 (width, height)
    :return: 滤波后的最小值
    """
    x, y = pos
    width, height = size
    rows, cols = len(data), len(data[0])
    min_val = float('inf')
    for i in range(max(0, y - height // 2), min(rows, y + height // 2 + 1)):
        for j in range(max(0, x - width // 2), min(cols, x + width // 2 + 1)):
            min_val = min(min_val, data[i][j])
    return min_val


def guidedFilter(guide, src, radius, eps):
    """
    引导滤波函数
    :param guide: 引导图像 (PIL Image)
    :param src: 输入图像 (PIL Image)
    :param radius: 滤波半径
    :param eps: 正则化参数
    :return: 滤波后的图像 (numpy array)
    """
    guide = np.array(guide, dtype=np.float64)
    src = np.array(src, dtype=np.float64)
    height, width = guide.shape

    # 计算局部均值
    window_size = 2 * radius + 1
    guide_mean = np.zeros_like(guide)
    src_mean = np.zeros_like(src)
    guide_squared_mean = np.zeros_like(guide)
    guide_src_mean = np.zeros_like(guide)

    for i in range(height):
        for j in range(width):
            i_min = max(0, i - radius)
            i_max = min(height, i + radius + 1)
            j_min = max(0, j - radius)
            j_max = min(width, j + radius + 1)
            guide_patch = guide[i_min:i_max, j_min:j_max]
            src_patch = src[i_min:i_max, j_min:j_max]
            guide_mean[i, j] = np.mean(guide_patch)
            src_mean[i, j] = np.mean(src_patch)
            guide_squared_mean[i, j] = np.mean(guide_patch ** 2)
            guide_src_mean[i, j] = np.mean(guide_patch * src_patch)

    # 计算局部方差和协方差
    guide_var = guide_squared_mean - guide_mean ** 2
    guide_src_cov = guide_src_mean - guide_mean * src_mean

    # 计算线性系数
    a = guide_src_cov / (guide_var + eps)
    b = src_mean - a * guide_mean

    # 计算系数的均值
    a_mean = np.zeros_like(a)
    b_mean = np.zeros_like(b)
    for i in range(height):
        for j in range(width):
            i_min = max(0, i - radius)
            i_max = min(height, i + radius + 1)
            j_min = max(0, j - radius)
            j_max = min(width, j + radius + 1)
            a_patch = a[i_min:i_max, j_min:j_max]
            b_patch = b[i_min:i_max, j_min:j_max]
            a_mean[i, j] = np.mean(a_patch)
            b_mean[i, j] = np.mean(b_patch)

    # 计算输出图像
    output = a_mean * guide + b_mean
    return output


def run_process():
    input_image_path = input_path_entry.get()
    output_image_path = output_path_entry.get()
    cut = float(cut_entry.get())
    omiga = float(omiga_entry.get())
    t0 = float(t0_entry.get())

    try:
        image = Image.open(input_image_path)
        image = image.convert('RGB')
        gray_image = image.convert('L')  # 转为灰度图以便进行区域生长
        # 区域生长分割天空
        seed_point = (image.size[0] // 2, 0)  # 假设种子点在图像顶部中央
        sky_mask = region_growing(gray_image, seed_point, threshold=30)  # 分割天空区域

        # 保存天空分割结果
        sky_mask_image = Image.fromarray((sky_mask * 255).astype(np.uint8))
        sky_mask_image.save('walker_sky_mask2.png')

        dark = getDark(image, minimizeFilter, (10, 10))
        dark.save('walker_dark2.png')

        light = getLight(image, dark, cut)

        transmission = getTransmission(image, light, omiga)
        # 对天空区域进行特殊处理
        sky_mask = np.array(sky_mask)  # 转换为 numpy array
        transmission = process_sky_region(image, sky_mask)

        tranImage = Image.new('L', image.size)
        for x in range(image.size[0]):
            for y in range(image.size[1]):
                tranImage.putpixel((x, y), int(transmission[y][x] * 255))

        guided = guidedFilter(gray_image, tranImage, 25, 0.001)

        guidedImage = Image.new('L', image.size)

        for x in range(image.size[0]):
            for y in range(image.size[1]):
                guidedImage.putpixel((x, y), ensure(guided[y][x]))
                guided[y][x] /= 255.0

        # guidedImage.show()
        guidedImage.save('walker_guided2.png')

        output = getRadiance(image, guided, light, t0)
        output.save(output_image_path)
        status_label.config(text="处理完成！")
    except Exception as e:
        status_label.config(text=f"处理出错: {str(e)}")


root = tk.Tk()
root.title("图像去雾处理")

# 输入图像路径
input_path_label = tk.Label(root, text="输入图像路径:")
input_path_label.pack()
input_path_entry = tk.Entry(root, width=50)
input_path_entry.pack()
input_path_button = tk.Button(root, text="选择文件",
                              command=lambda: input_path_entry.insert(0, filedialog.askopenfilename()))
input_path_button.pack()

# 输出图像路径
output_path_label = tk.Label(root, text="输出图像路径:")
output_path_label.pack()
output_path_entry = tk.Entry(root, width=50)
output_path_entry.pack()
output_path_button = tk.Button(root, text="选择文件",
                               command=lambda: output_path_entry.insert(0, filedialog.asksaveasfilename()))
output_path_button.pack()

# 其他参数输入
cut_label = tk.Label(root, text="Cut 参数:")
cut_label.pack()
cut_entry = tk.Entry(root)
cut_entry.insert(0, "0.001")
cut_entry.pack()

omiga_label = tk.Label(root, text="Omiga 参数:")
omiga_label.pack()
omiga_entry = tk.Entry(root)
omiga_entry.insert(0, "0.9")
omiga_entry.pack()

t0_label = tk.Label(root, text="t0 参数:")
t0_label.pack()
t0_entry = tk.Entry(root)
t0_entry.insert(0, "0.1")
t0_entry.pack()

# 运行按钮
run_button = tk.Button(root, text="运行处理", command=run_process)
run_button.pack()

# 状态显示
status_label = tk.Label(root, text="")
status_label.pack()

root.mainloop()